"""
文本处理工具函数
提供文本预处理、提示词优化等功能
"""

import re
import logging
from typing import List, Dict, Any, Optional
from pathlib import Path
import json

logger = logging.getLogger(__name__)

class TextProcessor:
    """文本处理器"""
    
    def __init__(self):
        """初始化文本处理器"""
        self.prompt_templates = self._load_prompt_templates()
        self.style_keywords = self._load_style_keywords()
    
    def _load_prompt_templates(self) -> Dict[str, str]:
        """加载提示词模板"""
        return {
            "animal": "一只{animal}在{scene}中{action}",
            "person": "一个{person}在{scene}中{action}",
            "scene": "在{scene}中，{description}",
            "action": "{subject}正在{action}",
            "emotion": "一个{emotion}的{subject}在{scene}中"
        }
    
    def _load_style_keywords(self) -> Dict[str, List[str]]:
        """加载风格关键词"""
        return {
            "realistic": ["realistic", "photorealistic", "high quality", "detailed", "sharp"],
            "cartoon": ["cartoon", "animated", "colorful", "bright", "fun"],
            "anime": ["anime", "Japanese animation", "manga style", "detailed"],
            "watercolor": ["watercolor", "painting", "artistic", "soft", "flowing"],
            "oil_painting": ["oil painting", "artistic", "textured", "brush strokes"],
            "sketch": ["sketch", "drawing", "black and white", "line art", "minimalist"]
        }
    
    def process(self, text: str) -> str:
        """处理文本提示"""
        # 清理文本
        text = self._clean_text(text)
        
        # 增强提示词
        text = self._enhance_prompt(text)
        
        # 添加质量关键词
        text = self._add_quality_keywords(text)
        
        return text
    
    def _clean_text(self, text: str) -> str:
        """清理文本"""
        # 移除多余空格
        text = re.sub(r'\s+', ' ', text.strip())
        
        # 移除特殊字符
        text = re.sub(r'[^\w\s\u4e00-\u9fff,，.。!！?？]', '', text)
        
        return text
    
    def _enhance_prompt(self, text: str) -> str:
        """增强提示词"""
        # 检测语言
        if self._is_chinese(text):
            text = self._enhance_chinese_prompt(text)
        else:
            text = self._enhance_english_prompt(text)
        
        return text
    
    def _is_chinese(self, text: str) -> bool:
        """检测是否为中文"""
        chinese_chars = re.findall(r'[\u4e00-\u9fff]', text)
        return len(chinese_chars) > len(text) * 0.5
    
    def _enhance_chinese_prompt(self, text: str) -> str:
        """增强中文提示词"""
        # 添加动作描述
        action_keywords = ["玩耍", "奔跑", "跳舞", "飞行", "游泳", "跳跃"]
        if not any(keyword in text for keyword in action_keywords):
            text += "，动作自然流畅"
        
        # 添加环境描述
        if "在" not in text and "里" not in text and "中" not in text:
            text += "，背景优美"
        
        # 添加质量描述
        quality_keywords = ["高清", "清晰", "精美", "高质量"]
        if not any(keyword in text for keyword in quality_keywords):
            text += "，画面精美"
        
        return text
    
    def _enhance_english_prompt(self, text: str) -> str:
        """增强英文提示词"""
        # 添加动作描述
        action_keywords = ["playing", "running", "dancing", "flying", "swimming", "jumping"]
        if not any(keyword in text.lower() for keyword in action_keywords):
            text += ", natural movement"
        
        # 添加环境描述
        if "in" not in text.lower() and "at" not in text.lower():
            text += ", beautiful background"
        
        # 添加质量描述
        quality_keywords = ["high quality", "detailed", "sharp", "beautiful"]
        if not any(keyword in text.lower() for keyword in quality_keywords):
            text += ", high quality"
        
        return text
    
    def _add_quality_keywords(self, text: str) -> str:
        """添加质量关键词"""
        quality_suffixes = [
            ", high quality, detailed, sharp focus",
            ", beautiful composition, professional lighting",
            ", cinematic, smooth motion, fluid animation"
        ]
        
        # 随机选择一个质量后缀
        import random
        suffix = random.choice(quality_suffixes)
        
        return text + suffix
    
    def add_style_keywords(self, text: str, style: str) -> str:
        """添加风格关键词"""
        if style in self.style_keywords:
            keywords = self.style_keywords[style]
            style_text = ", ".join(keywords)
            return f"{text}, {style_text}"
        
        return text
    
    def generate_prompt_from_template(self, template_name: str, **kwargs) -> str:
        """从模板生成提示词"""
        if template_name in self.prompt_templates:
            template = self.prompt_templates[template_name]
            try:
                return template.format(**kwargs)
            except KeyError as e:
                logger.warning(f"模板参数缺失: {e}")
                return template
        
        return ""
    
    def extract_keywords(self, text: str) -> List[str]:
        """提取关键词"""
        # 简单的关键词提取
        words = text.split()
        
        # 过滤停用词
        stop_words = {"的", "了", "在", "是", "有", "和", "与", "或", "the", "a", "an", "and", "or", "in", "on", "at", "to", "for", "of", "with", "by"}
        
        keywords = [word for word in words if word.lower() not in stop_words and len(word) > 1]
        
        return keywords[:10]  # 返回前10个关键词
    
    def analyze_prompt(self, text: str) -> Dict[str, Any]:
        """分析提示词"""
        analysis = {
            "length": len(text),
            "language": "chinese" if self._is_chinese(text) else "english",
            "keywords": self.extract_keywords(text),
            "has_action": self._has_action_keywords(text),
            "has_emotion": self._has_emotion_keywords(text),
            "has_style": self._has_style_keywords(text)
        }
        
        return analysis
    
    def _has_action_keywords(self, text: str) -> bool:
        """检查是否包含动作关键词"""
        action_keywords = ["玩耍", "奔跑", "跳舞", "飞行", "游泳", "跳跃", "playing", "running", "dancing", "flying", "swimming", "jumping"]
        return any(keyword in text.lower() for keyword in action_keywords)
    
    def _has_emotion_keywords(self, text: str) -> bool:
        """检查是否包含情感关键词"""
        emotion_keywords = ["开心", "快乐", "悲伤", "愤怒", "惊讶", "happy", "sad", "angry", "surprised", "excited"]
        return any(keyword in text.lower() for keyword in emotion_keywords)
    
    def _has_style_keywords(self, text: str) -> bool:
        """检查是否包含风格关键词"""
        all_style_keywords = []
        for keywords in self.style_keywords.values():
            all_style_keywords.extend(keywords)
        
        return any(keyword in text.lower() for keyword in all_style_keywords)
    
    def optimize_prompt(self, text: str, target_length: int = 100) -> str:
        """优化提示词长度"""
        if len(text) <= target_length:
            return text
        
        # 分割句子
        sentences = re.split(r'[，。！？,\.!?]', text)
        
        # 保留最重要的句子
        optimized_sentences = []
        current_length = 0
        
        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue
            
            if current_length + len(sentence) <= target_length:
                optimized_sentences.append(sentence)
                current_length += len(sentence)
            else:
                break
        
        return "，".join(optimized_sentences) + "。"
    
    def create_variations(self, text: str, num_variations: int = 3) -> List[str]:
        """创建提示词变体"""
        variations = []
        
        # 基础变体
        base_variations = [
            text,
            text.replace("，", "，同时"),
            text.replace("。", "，画面精美。"),
            f"在{text}中",
            f"场景：{text}"
        ]
        
        # 添加风格变体
        for style in ["realistic", "cartoon", "anime"]:
            style_text = self.add_style_keywords(text, style)
            base_variations.append(style_text)
        
        # 选择指定数量的变体
        import random
        variations = random.sample(base_variations, min(num_variations, len(base_variations)))
        
        return variations
    
    def translate_prompt(self, text: str, target_language: str = "english") -> str:
        """翻译提示词"""
        # 简单的翻译映射
        translations = {
            "一只可爱的小猫": "a cute cat",
            "在花园里玩耍": "playing in the garden",
            "小女孩": "little girl",
            "跳舞": "dancing",
            "花海": "flower field",
            "宇航员": "astronaut",
            "太空": "space",
            "漂浮": "floating",
            "舞者": "dancer",
            "舞台": "stage",
            "芭蕾": "ballet",
            "日落": "sunset",
            "海滩": "beach",
            "波浪": "waves",
            "森林": "forest",
            "瀑布": "waterfall",
            "流水": "flowing water",
            "城市": "city",
            "夜景": "night scene",
            "车流": "traffic"
        }
        
        if target_language == "english":
            for chinese, english in translations.items():
                text = text.replace(chinese, english)
        
        return text
    
    def save_prompt_history(self, prompt: str, output_path: str, metadata: Dict[str, Any] = None):
        """保存提示词历史"""
        history_file = Path("outputs/prompt_history.jsonl")
        history_file.parent.mkdir(parents=True, exist_ok=True)
        
        entry = {
            "prompt": prompt,
            "timestamp": time.time(),
            "output_path": output_path,
            "metadata": metadata or {}
        }
        
        with open(history_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
    
    def load_prompt_history(self, limit: int = 100) -> List[Dict[str, Any]]:
        """加载提示词历史"""
        history_file = Path("outputs/prompt_history.jsonl")
        
        if not history_file.exists():
            return []
        
        history = []
        with open(history_file, "r", encoding="utf-8") as f:
            for line in f:
                if len(history) >= limit:
                    break
                try:
                    entry = json.loads(line.strip())
                    history.append(entry)
                except json.JSONDecodeError:
                    continue
        
        return history

def clean_text(text: str) -> str:
    """清理文本"""
    processor = TextProcessor()
    return processor._clean_text(text)

def enhance_prompt(text: str) -> str:
    """增强提示词"""
    processor = TextProcessor()
    return processor.process(text)

def extract_keywords(text: str) -> List[str]:
    """提取关键词"""
    processor = TextProcessor()
    return processor.extract_keywords(text)

def analyze_prompt(text: str) -> Dict[str, Any]:
    """分析提示词"""
    processor = TextProcessor()
    return processor.analyze_prompt(text) 